import sys
import time
from pathlib import Path
import threading

import numpy as np

try:
    from primesense import openni2, _openni2 as c_api  # type: ignore
    HAS_OPENNI = True
except Exception:
    HAS_OPENNI = False

try:
    import cv2
    HAS_CV = True
except Exception:
    HAS_CV = False


def try_init_openni() -> str | None:
    if not HAS_OPENNI:
        return None
    # 1) default
    try:
        openni2.initialize()
        return "default"
    except Exception:
        pass
    # 2) common install paths
    candidates = [
        r"C:\\Program Files\\OpenNI2\\Redist",
        r"C:\\Program Files (x86)\\OpenNI2\\Redist",
        r"C:\\Program Files\\Orbbec\\OpenNI2\\Redist",
        r"C:\\Program Files (x86)\\Orbbec\\OpenNI2\\Redist",
    ]
    for p in candidates:
        try:
            openni2.initialize(p)
            return p
        except Exception:
            continue
    return None


def diagnose_openni() -> None:
    print("[OpenNI2] Python wrapper:", "OK" if HAS_OPENNI else "MISSING")
    if not HAS_OPENNI:
        print("[OpenNI2] 请安装 primesense 包，并确保已安装 OpenNI2/Orbbec SDK")
        return

    path = try_init_openni()
    print("[OpenNI2] initialize:", path or "FAILED")
    if path is None:
        print("[OpenNI2] 初始化失败，可能未安装 OpenNI2/Orbbec SDK 或 PATH 未包含 Redist")
        return

    dev = None
    try:
        dev = openni2.Device.open_any()
        print("[OpenNI2] open_any: OK")
    except Exception as e:
        # 尝试 enumerate_uris
        enum_uris = getattr(openni2.Device, 'enumerate_uris', None)
        if enum_uris is not None:
            try:
                uris = enum_uris()
                print("[OpenNI2] enumerate_uris:", uris)
                if uris:
                    dev = openni2.Device.open(uris[0])
                    print("[OpenNI2] open by uri: OK")
            except Exception as e2:
                print("[OpenNI2] open failed:", e, e2)
        else:
            print("[OpenNI2] open_any failed:", e)
    if dev is None:
        openni2.unload()
        return

    # Device info
    force_skip_color = False
    try:
        info = dev.get_device_info()
        print(f"[OpenNI2] Device Name: {info.name}")
        print(f"[OpenNI2] Vendor: {info.vendor}")
        print(f"[OpenNI2] USB Product ID: {info.usbProductId}")
        print(f"[OpenNI2] USB Vendor ID: {info.usbVendorId}")
        
        # Astra Pro (PID 1027) 的 OpenNI2 彩色流会在 start() 时卡死
        if info.usbProductId == 1027:
            force_skip_color = True
            print(f"[OpenNI2] ⚠️ 检测到 Astra Pro (PID=1027)，已知 color.start() 会卡死，将跳过测试")
    except Exception as e:
        print(f"[OpenNI2] get_device_info failed: {e}")

    # Registration support
    try:
        dev.set_image_registration_mode(openni2.IMAGE_REGISTRATION_DEPTH_TO_COLOR)
        print("[OpenNI2] registration: DEPTH_TO_COLOR")
    except Exception:
        print("[OpenNI2] registration: not supported")

    # Depth stream - 列出所有支持的分辨率
    depth = None
    try:
        depth = dev.create_depth_stream()
        print("[OpenNI2] Depth stream created: OK")
        
        # 列出所有支持的视频模式
        try:
            video_modes = depth.get_supported_video_modes()
            print(f"[OpenNI2] Supported depth modes ({len(video_modes)} total):")
            for i, mode in enumerate(video_modes[:10]):  # 只显示前10个
                print(f"  [{i}] {mode.resolutionX}x{mode.resolutionY} @ {mode.fps}fps, fmt={mode.pixelFormat}")
        except Exception as e:
            print(f"[OpenNI2] get_supported_video_modes failed: {e}")
        
        dvm = c_api.OniVideoMode()
        dvm.pixelFormat = c_api.OniPixelFormat.ONI_PIXEL_FORMAT_DEPTH_1_MM
        dvm.resolutionX = 640
        dvm.resolutionY = 480
        dvm.fps = 30
        depth.set_video_mode(dvm)
        depth.start()
        dframe = depth.read_frame()
        dnp = np.frombuffer(dframe.get_buffer_as_uint16(), dtype=np.uint16)
        dnp = dnp.reshape((480, 640))
        valid_pixels = np.count_nonzero(dnp > 0)
        print(f"[OpenNI2] depth: {dvm.resolutionX}x{dvm.resolutionY}@{dvm.fps}fps  valid_pixels={valid_pixels}/{dnp.size}")
    except Exception as e:
        print("[OpenNI2] depth stream failed:", e)
        depth = None

    # Color stream (RGB888 -> fallback YUYV)
    color = None
    color_fmt = None
    
    if force_skip_color:
        print("[OpenNI2] Color stream test: SKIPPED (Astra Pro 已知问题)")
        print("[OpenNI2] 提示: Astra Pro 需要使用 UVC 作为彩色流")
    else:
        try:
            color = dev.create_color_stream()
            print("[OpenNI2] Color stream created: OK")
            
            # 列出所有支持的彩色模式
            try:
                color_modes = color.get_supported_video_modes()
                print(f"[OpenNI2] Supported color modes ({len(color_modes)} total):")
                for i, mode in enumerate(color_modes[:10]):  # 只显示前10个
                    print(f"  [{i}] {mode.resolutionX}x{mode.resolutionY} @ {mode.fps}fps, fmt={mode.pixelFormat}")
            except Exception as e:
                print(f"[OpenNI2] get_supported_video_modes (color) failed: {e}")
            
            cvm = c_api.OniVideoMode()
            cvm.resolutionX = 640
            cvm.resolutionY = 480
            cvm.fps = 30
            try:
                cvm.pixelFormat = c_api.OniPixelFormat.ONI_PIXEL_FORMAT_RGB888
                color.set_video_mode(cvm)
                color_fmt = 'RGB888'
                print("[OpenNI2] Color format: RGB888")
            except Exception as e1:
                print(f"[OpenNI2] RGB888 set_video_mode failed: {e1}, trying YUYV...")
                try:
                    cvm.pixelFormat = c_api.OniPixelFormat.ONI_PIXEL_FORMAT_YUYV
                    color.set_video_mode(cvm)
                    color_fmt = 'YUYV'
                    print("[OpenNI2] Color format: YUYV")
                except Exception as e2:
                    print(f"[OpenNI2] YUYV set_video_mode failed: {e2}")
                    raise
            
            # 对于其他设备，尝试启动（但这部分不会执行到 Astra Pro）
            color.start()
            cframe = color.read_frame()
            cbuf = np.frombuffer(cframe.get_buffer_as_uint8(), dtype=np.uint8)
            print(f"[OpenNI2] color: {cvm.resolutionX}x{cvm.resolutionY}@{cvm.fps}fps  fmt={color_fmt}  buffer_size={cbuf.size}")
        except Exception as e:
            print("[OpenNI2] color stream TOTAL FAIL:", e)
            print("[OpenNI2] 提示: 彩色流不支持 OpenNI2，需要 UVC 回退")

    try:
        if color is not None:
            color.stop()
        if depth is not None:
            depth.stop()
    except Exception:
        pass
    openni2.unload()


def diagnose_uvc() -> None:
    print("[UVC] OpenCV availability:", "OK" if HAS_CV else "MISSING")
    if not HAS_CV:
        print("[UVC] 请安装 opencv-python")
        return
    found = False
    for idx in range(0, 3):
        cap = cv2.VideoCapture(idx, cv2.CAP_DSHOW)
        ok = cap.isOpened()
        ret, frame = cap.read()
        shape = frame.shape if ret else None
        print(f"[UVC] index={idx} opened={ok} read={ret} shape={shape}")
        if ok and ret:
            found = True
        cap.release()
    if not found:
        print("[UVC] 未检测到可读的UVC彩色视频源，可能被占用或权限受限")


def main() -> None:
    print("=== Camera/Driver Diagnostic ===")
    diagnose_openni()
    print("--------------------------------")
    diagnose_uvc()
    print("--------------------------------")
    print("完成。如 OpenNI2 或 UVC 任何一项失败，请重插USB3口、检查相机权限、关闭占用相机的软件，并确保已安装 Orbbec/OpenNI2 SDK。")


if __name__ == '__main__':
    main()


